{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Functions and Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy import (\n",
    "    asarray,\n",
    "    corrcoef,\n",
    "    diag,\n",
    "    eye,\n",
    "    hstack,\n",
    "    logical_and,\n",
    "    logical_not,\n",
    "    r_,\n",
    "    sort,\n",
    "    unique,\n",
    "    zeros,\n",
    ")\n",
    "from numpy.linalg import inv\n",
    "from numpy.random import chisquare, standard_normal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ascategory(x):\n",
    "    \"\"\" \"\"\"\n",
    "    x = asarray(x)\n",
    "    t = x.shape[0]\n",
    "    ux = unique(x)\n",
    "    k = ux.shape[0]\n",
    "    categories = np.zeros((t, k))\n",
    "    for i in range(k):\n",
    "        loc = np.squeeze(x == ux[i])\n",
    "        categories[loc, i] = 1.0\n",
    "    return categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ascategory([\"a\", \"b\", \"a\", \"c\", \"a\", \"b\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gls(x, y, omega=None):\n",
    "    \"\"\" \"\"\"\n",
    "    t, k = x.shape\n",
    "    if omega is None:\n",
    "        omega = eye(t)\n",
    "    omega_inv = inv(omega)\n",
    "    xpx = x.T @ omega_inv @ x\n",
    "    xpy = x.T @ omega_inv @ y\n",
    "    beta_gls = asarray(inv(xpx) @ xpy)\n",
    "    return beta_gls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "x = standard_normal((100, 3))\n",
    "y = standard_normal((100, 1))\n",
    "gls(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "omega = np.diag(chisquare(5, size=100))\n",
    "omega"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "gls(x, y, omega)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def partial_corr(x, y=None, quantile=0.5, tail=\"Lower\"):\n",
    "    \"\"\" \"\"\"\n",
    "    if y is not None:\n",
    "        X = x.view()\n",
    "        Y = y.view()\n",
    "        T = X.shape[0]\n",
    "        X.shape = T, 1\n",
    "        Y.shape = T, 1\n",
    "        z = hstack((X, Y))\n",
    "    else:\n",
    "        z = x\n",
    "    T, K = z.shape\n",
    "    corr = eye(K)\n",
    "    count = zeros((K, K))\n",
    "    ind = zeros((T, K), dtype=np.bool)\n",
    "    for i in range(K):\n",
    "        temp = sort(z[:, i].ravel())\n",
    "        cutoff = int(round(quantile * T))\n",
    "        threshold = temp[cutoff]\n",
    "        ind[:, i] = z[:, i] < threshold\n",
    "        if tail == \"Upper\":\n",
    "            ind[:, i] = logical_not(ind[:, i])\n",
    "    for i in range(K):\n",
    "        for j in range(i + 1, K):\n",
    "            pl = logical_and(ind[:, i], ind[:, j])\n",
    "            count[i, j] = sum(pl)\n",
    "            count[j, i] = count[i, j]\n",
    "            if sum(pl) > 1:\n",
    "                w = z[pl, :]\n",
    "                w = w[:, r_[i, j]]\n",
    "                corr[i, j] = corrcoef(w.T)[0, 1]\n",
    "                corr[j, i] = corr[i, j]\n",
    "            else:\n",
    "                corr[i, j] = np.nan\n",
    "                corr[j, i] = np.nan\n",
    "    return corr, count"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
